#include "stdafx.h"
#include "Common.h"
#include "WMP.h"
#include "WMPHook.h"
#include "User32Hook.h"
#include "Shell32Hook.h"
#include "OLE32Hook.h"
#include "WMVCoreHook.h"
#include "MSVCRTHook.h"
#include "PropSysHook.h"
#include "MFHook.h"
#include "Patch.h"
#include "Kernel32Hook.h"
#include "QuartzHook.h"

#pragma unmanaged

LPCSTR kernel32Methods[] = {
	"FreeLibrary",
	"GetProcAddress",
	"GetModuleHandleA",
	"GetModuleHandleW"
};
const int kernel32Count = sizeof(kernel32Methods) / sizeof(kernel32Methods[0]);

LPCSTR user32Methods[] = {
	"TrackPopupMenu",
	"CallWindowProcW",
	"LoadMenuW",
	"DialogBoxParamW",
	"DispatchMessageW",
#ifdef _WIN64
	"SetWindowLongPtrW"
#else
	"SetWindowLongW"
#endif
};
const int user32Count = sizeof(user32Methods) / sizeof(user32Methods[0]);

LPCSTR mfMethods[] = {
	"MFCreateMediaSession",
	"MFCreatePMPMediaSession"
};
const int mfCount = sizeof(mfMethods) / sizeof(mfMethods[0]);

LPCSTR msvcrtMethods[] = {
	"_wcsicmp"
};
const int msvcrtCount = sizeof(msvcrtMethods) / sizeof(msvcrtMethods[0]);

LPCSTR ole32Methods[] = {
	"CoCreateInstance"
};
const int ole32Count = sizeof(ole32Methods) / sizeof(ole32Methods[0]);

LPCSTR apilibraryloaderMethods[] = { 
	"GetProcAddress"
};
const int apilibraryloaderCount = sizeof(apilibraryloaderMethods) / sizeof(apilibraryloaderMethods[0]);

LPCSTR propsysMethods[] = {
	"PSLookupPropertyHandlerCLSID"
};
const int propsysCount = sizeof(propsysMethods) / sizeof(propsysMethods[0]);

LPCSTR shell32Methods[] = {
	"SHOpenFolderAndSelectItems"
};
const int shell32Count = sizeof(shell32Methods) / sizeof(shell32Methods[0]);

LPCSTR wmvcoreMethods[] = {
	"WMCreateEditor",
	"WMCheckURLExtension",
	"WMCreateReader"
};
const int wmvcoreCount = sizeof(wmvcoreMethods) / sizeof(wmvcoreMethods[0]);

const int wmpOrdinals[] = {
	3000,
	3001,
	4000,
	3002,
	3003,
	3004,
	3005,
	0, 0, 0, 0, 0
};
LPCSTR wmpMethods[] = {
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	"DllCanUnloadNow",
	"DllGetClassObject",
	"DllRegisterServer",
	"DllUnregisterServer",
	"InitializeTests"
};
const int wmpCount = sizeof(wmpMethods) / sizeof(wmpMethods[0]);

LPCSTR Kernel32Methods[] = {
	"CreateThreadpoolTimer",
	"CloseThreadpoolTimer",
	"GetTickCount"
};
const int Kernel32Count = sizeof(Kernel32Methods) / sizeof(Kernel32Methods[0]);

LPCSTR classObjectMethods[] = {
	"DllGetClassObject"
};
const int classObjectCount = sizeof(classObjectMethods) / sizeof(classObjectMethods[0]);

LPCSTR quartzExternals[] = {
	"QuartzGetClassObject"
};


void* user32Save[user32Count] = {NULL, NULL, NULL, NULL, NULL, NULL};
void* user32Patch[user32Count] = {TrackPopupMenuHook, CallWindowProcWHook, LoadMenuWHook, DialogBoxParamWHook, DispatchMessageWHook, SetWindowLongPtrWHook};
void* user32Hook[user32Count] = {NULL, NULL, NULL, NULL, NULL, NULL};
void* shell32Save[shell32Count] = {NULL};
void* shell32Patch[shell32Count] = {SHOpenFolderAndSelectItemsHook};
void* shell32Hook[shell32Count] = {NULL};
void* ole32Save[ole32Count] = {NULL};
void* ole32Patch[ole32Count] = {CoCreateInstanceHook};
void* ole32Hook[ole32Count] = {NULL};
void* wmvcoreSave[wmvcoreCount] = {NULL, NULL, NULL};
void* wmvcorePatch[wmvcoreCount] = {WMCreateEditorHook, WMCheckURLExtensionHook, WMCreateReaderHook};
void* wmvcoreHook[wmvcoreCount] = {NULL, NULL, NULL};
void* msvcrtSave[msvcrtCount] = {NULL};
void* msvcrtPatch[msvcrtCount] = {wcsicmpHook};
void* msvcrtHook[msvcrtCount] = {NULL};
void* propsysSave[propsysCount] = {NULL};
void* propsysPatch[propsysCount] = {PSLookupPropertyHandlerCLSIDHook};
void* propsysHook[propsysCount] = {NULL};
void* mfSave[mfCount] = {NULL, NULL};
void* mfPatch[mfCount] = {MFCreateMediaSessionHook, MFCreatePMPMediaSessionHook};
void* mfHook[mfCount] = {NULL, NULL};
void* Kernel32Save[Kernel32Count] = {NULL, NULL, NULL};
void* Kernel32Patch[Kernel32Count] = {CreateThreadpoolTimerHook, CloseThreadpoolTimerHook, GetTickCountHook};
void* Kernel32Hook[Kernel32Count] = {NULL, NULL, NULL};
void* quartzSave[classObjectCount] = {NULL};
void* quartzPatch[classObjectCount] = {QuartzGetClassObjectHook};
void* quartzHook[classObjectCount] = {NULL};

void* wmpSave[wmpCount] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
void* kernel32Patch[kernel32Count] = {FreeLibraryHook, GetProcAddressHook, GetModuleHandleAHook, GetModuleHandleWHook};
void* kernel32Save[kernel32Count] = {NULL, NULL, NULL, NULL};
void* apilibraryloaderSave[apilibraryloaderCount] = {NULL};
void* apilibraryloaderPatch[apilibraryloaderCount] = {LoaderGetProcAddressHook};

typedef BOOL (WINAPI *PFreeLibrary)(HMODULE hLibModule);
typedef FARPROC (WINAPI *PGetProcAddress)(HMODULE hModule, LPCSTR lpProcName);
typedef HMODULE (WINAPI *PGetModuleHandleA)(LPCSTR lpModuleName);
typedef HMODULE (WINAPI *PGetModuleHandleW)(LPCWSTR lpModuleName);

void* FixPushCommand(void* handle) {
	try {
		if (!handle || *(char*)handle != 0x68)
			return handle;
		//read PUSH operand
		return *(void**)((char*)handle + 1);
	} catch(...) {
		return handle;
	}
}

void SetPrevFuncAddress(LPCSTR method, void* proc) {
	CSLocker->AcquireLockShared();
	try {
		if (RestoreProcs) {
			for(int k = 0; k < RestoreProcs->GetCount(); k++) {
				if (lstrcmpiA(method, RestoreProcs->Get(k)) == 0) {
					PSetPrevAddress restore = RestoreProcs->GetData(k);
					if (restore)
						(*restore)(method, proc);
					break;
				}
			}
		}
	} catch (...) {
		TraceException(_T("SetPrevFuncAddress"));
	}
	CSLocker->ReleaseLockShared();
}

bool PatchLibrary(void* procs[], int count, void* patch[], void* save[], HMODULE handle, PIMAGE_IMPORT_DESCRIPTOR table) {
	bool patched = false;
	for (void** entry = (void**)((char*)handle + table->FirstThunk); *entry; entry++) {
		for (int k = 0; k < count; k++) {
			if (procs[k] != NULL && FixPushCommand(*entry) == procs[k] && patch[k])
			{
				DWORD flags;
				VirtualProtect(entry, 4, PAGE_EXECUTE_READWRITE, &flags);
				try {
					save[k] = *entry;
					*entry = patch[k];
					patched = true;
				}
				finally (
					VirtualProtect(entry, 4, flags, &flags);
				)
			}
		}
	}
	return patched;
}

void ReadMethods(void* procs[], LPCTSTR module, const char* methods[], int count) {
	HMODULE handle = GetModuleHandle(module);
	for (int k = 0; k < count; k++) {
		void* proc = GetProcAddress(handle, methods[k]);
		procs[k] = FixPushCommand(proc);
	}
}

int PatchModule(HMODULE handle, PATCH_ACTION action, IMPORT_LIBRARY* patched) {
	int hr = 0x0A;
	void* kernel32Procs[kernel32Count];
	void* Kernel32Procs[Kernel32Count];
	void* user32Procs[user32Count];
	void* msvcrtProcs[msvcrtCount];
	void* ole32Procs[ole32Count];
	void* mfProcs[mfCount];
	void* apilibraryloaderProcs[apilibraryloaderCount];
	int* result = (int*)patched;
	CreateTextBuffer(textBuffer);
	CreatePathBuffer(pathBuffer);
	*result = 0;

	try {
		if (!handle)
			return hr;
		hr = 0xEE;
		if (HInstance == handle)
			return hr;
		hr = 0xF5;
		PIMAGE_DOS_HEADER dosheader = (PIMAGE_DOS_HEADER)handle;
		PIMAGE_NT_HEADERS header = (PIMAGE_NT_HEADERS)((char*)handle + dosheader->e_lfanew);
		hr = 0xF4;
		PIMAGE_IMPORT_DESCRIPTOR importlib =  (PIMAGE_IMPORT_DESCRIPTOR)((char*)handle +
			header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
		hr = 0xF2;
		ReadMethods(kernel32Procs, KERNEL32DLL, kernel32Methods, kernel32Count);
		switch (action) {
		case PATCH_ACTION_WMPLAYER: {
				GetModuleFileName(NULL, pathBuffer, pathBufferLength);
				_sntprintf(textBuffer, textBufferLength, _T("Patching %s\r\n"), pathBuffer);
				OutputDebugString(textBuffer);
				break;
			}
		case PATCH_ACTION_WMP: {
				ReadMethods(Kernel32Procs, KERNEL32DLL, Kernel32Methods, Kernel32Count);
				ReadMethods(user32Procs, USER32DLL, user32Methods, user32Count);
				ReadMethods(msvcrtProcs, MSVCRTDLL, msvcrtMethods, msvcrtCount);
				ReadMethods(ole32Procs, OLE32DLL, ole32Methods, ole32Count);
				ReadMethods(mfProcs, MFDLL, mfMethods, mfCount);
				break;
			}
		case PATCH_ACTION_SHELL32: {
				ReadMethods(apilibraryloaderProcs, APILIBRARYLOADERDLL, apilibraryloaderMethods, apilibraryloaderCount);
				break;
			}
		case PATCH_ACTION_OLE32: {
				ReadMethods(apilibraryloaderProcs, APILIBRARYLOADERDLL, apilibraryloaderMethods, apilibraryloaderCount);
				break;
			}
		}
		hr = 0xF1;
		while (importlib->FirstThunk != 0) {
			hr = 0xF0;
			int offset = importlib->Name;
			LPSTR name = (LPSTR)((char*)handle + offset);
			hr = 0xEF;
			if ((action < PATCH_ACTION_2 || action == PATCH_ACTION_OLE32) && lstrcmpiA(KERNEL32DLLA, name) == 0 &&
				PatchLibrary(kernel32Procs, kernel32Count, kernel32Patch, kernel32Save, handle, importlib))
				*result |= IMPORT_LIBRARY_KERNEL32;
			switch (action) {
			case PATCH_ACTION_WMP: {
					if (lstrcmpiA(KERNEL32DLLA, name) == 0) {
						if (PatchLibrary(Kernel32Procs, Kernel32Count, Kernel32Patch, Kernel32Save, handle, importlib))
							*result |= IMPORT_LIBRARY_KERNEL32HOOK;
					} else if (lstrcmpiA(USER32DLLA, name) == 0) {
						if (PatchLibrary(user32Procs, user32Count, user32Patch, user32Save, handle, importlib))
							*result |= IMPORT_LIBRARY_USER32;
					} else if (lstrcmpiA(MSVCRTDLLA, name) == 0) {
						if (PatchLibrary(msvcrtProcs, msvcrtCount, msvcrtPatch, msvcrtSave, handle, importlib))
							*result |= IMPORT_LIBRARY_MSVCRT;
					} else if (lstrcmpiA(OLE32DLLA, name) == 0) {
						if (PatchLibrary(ole32Procs, ole32Count, ole32Patch, ole32Save, handle, importlib))
							*result |= IMPORT_LIBRARY_OLE32;
					} else if (lstrcmpiA(MFDLLA, name) == 0x0) {
						if (PatchLibrary(mfProcs, mfCount, mfPatch, mfSave, handle, importlib))
							*result |= IMPORT_LIBRARY_MF;
					}
					break;
				}
			case PATCH_ACTION_SHELL32: {
					if (lstrcmpiA(APILIBRARYLOADERDLLA, name) == 0) {
						if (PatchLibrary(apilibraryloaderProcs, apilibraryloaderCount, apilibraryloaderPatch, apilibraryloaderSave, handle, importlib))
							*result |= IMPORT_LIBRARY_API_LIBRARYLOADER;
					}
					break;
				}
			case PATCH_ACTION_OLE32: {
					if (lstrcmpiA(APILIBRARYLOADERDLLA, name) == 0) {
						if (PatchLibrary(apilibraryloaderProcs, apilibraryloaderCount, apilibraryloaderPatch, apilibraryloaderSave, handle, importlib))
							*result |= IMPORT_LIBRARY_API_LIBRARYLOADER;
					}
					break;
				}
			default: {
					break;
				}
			}
			importlib++;
		}
		hr = 0x0;
	} catch (...) {
		TraceException(_T("PatchModule"));
	}
	return hr;
}

void TraceAppendLibraryName(LPCTSTR name, LPTSTR textBuffer, int &k) {
	if (k) _tcscat(textBuffer, _T(", "));
	else _tcscat(textBuffer, _T(" ("));
	k++;
	_tcscat(textBuffer, name);
}

void TraceResult(LPCTSTR path, int hr, IMPORT_LIBRARY patched) {
	CreateTextBuffer(textBuffer);

	if (hr == 0) {
		_sntprintf(textBuffer, textBufferLength, _T("Hooking of %s : %d"), path, patched);
		int k = 0;
		if (patched & IMPORT_LIBRARY_KERNEL32) TraceAppendLibraryName(KERNEL32DLL _T(" (internal)"), textBuffer, k);
		if (patched & IMPORT_LIBRARY_USER32) TraceAppendLibraryName(USER32DLL, textBuffer, k);
		if (patched & IMPORT_LIBRARY_MSVCRT) TraceAppendLibraryName(MSVCRTDLL, textBuffer, k);
		if (patched & IMPORT_LIBRARY_KERNEL32HOOK) TraceAppendLibraryName(KERNEL32DLL, textBuffer, k);
		if (patched & IMPORT_LIBRARY_API_LIBRARYLOADER) TraceAppendLibraryName(APILIBRARYLOADERDLL, textBuffer, k);
		if (patched & IMPORT_LIBRARY_OLE32) TraceAppendLibraryName(OLE32DLL, textBuffer, k);
		if (patched & IMPORT_LIBRARY_MF) TraceAppendLibraryName(MFDLL, textBuffer, k);
		if (k) _tcscat(textBuffer, _T(")\r\n"));
		else _tcscat(textBuffer, _T("\r\n"));
	} else {
		_sntprintf(textBuffer,  textBufferLength, _T("Hooking of %s failed : %X\r\n"), path, hr);
	}
	OutputDebugString(textBuffer);
}

IMPORT_LIBRARY Patch(HMODULE handle, LPCTSTR path, PATCH_ACTION action) {
	IMPORT_LIBRARY lib;
	int hr = PatchModule(handle, action, &lib);
	TraceResult(path, hr, lib);
	return lib;
}

BOOL WINAPI FreeLibraryHook(HMODULE hLibModule) {
	CreatePathBuffer(pathBuffer);

	if (hLibModule == HInstance) {
		BOOL res = FALSE;
		if (HINST)
			res = (*(PFreeLibrary)kernel32Save[0])(HINST);
		if (GetModuleFileName(HINST, pathBuffer, pathBufferLength) == 0) {
			HINST = NULL;
			Uninitialize();
		}
		return res;
	} else {
		return (*(PFreeLibrary)kernel32Save[0])(hLibModule);
	}
}

void* HookMethod(LPCSTR method, void*** psave, void* patch[], LPCSTR methods[], void* save[], int count, LPCSTR* pexternal = NULL, LPCSTR externals[] = NULL) {
	if (((INT_PTR)method >> 0x10) != 0x0) {
		for (int k = 0; k < count; k++) {
			if (lstrcmpiA(method, methods[k]) == 0) {
				*psave = &save[k];
				if (pexternal && externals && externals[k])
					*pexternal = externals[k];
				return patch[k];
			}
		}
	}
	return NULL;
}

FARPROC WINAPI LoaderGetProcAddressHook(HMODULE hModule, LPCSTR lpProcName) {
	void* patch = NULL;
	void** saved = NULL;
	LPCSTR method = lpProcName;
	CreatePathBuffer(pathBuffer);
	CreatePathBuffer(nameBuffer);

	try {
		GetModuleFileName(hModule, pathBuffer, pathBufferLength);
		GetFileName(pathBuffer, nameBuffer);
		if (lstrcmpi(nameBuffer, PROPSYSDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, propsysPatch, propsysMethods, propsysSave, propsysCount);
		} else if (lstrcmpi(nameBuffer, QUARTZDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, quartzPatch, classObjectMethods, quartzSave, classObjectCount, &method, quartzExternals);
		}
	} catch (...) {
		TraceException(_T("NewLoaderGetProcAddress"));
	}
	void* proc = (*(PGetProcAddress)apilibraryloaderSave[0])(hModule, lpProcName);
	if (!patch || !saved) {
		patch = proc;
	} else {
		*saved = proc;
		if (proc)
			SetPrevFuncAddress(method, proc);
	}
	return (FARPROC)patch;
}

FARPROC WINAPI GetProcAddressHook(HMODULE hModule, LPCSTR lpProcName) {
	void* patch = NULL;
	bool iskernel32 = false;
	void** saved = NULL;
	LPCSTR method = lpProcName;
	CreatePathBuffer(pathBuffer);
	CreatePathBuffer(nameBuffer);

	if (HInstance == hModule) {
		Initialize(true);
		if (HINST)
			patch = (*(PGetProcAddress)kernel32Save[1])(HINST, lpProcName);
		return (FARPROC)patch;
	}
	if (!Initialized || Uninitialized) {
		patch = (*(PGetProcAddress)kernel32Save[1])(hModule, lpProcName);
		return (FARPROC)patch;
	}
	try {
		GetModuleFileName(hModule, pathBuffer, pathBufferLength);
		GetFileName(pathBuffer, nameBuffer);
		if (lstrcmpi(nameBuffer, KERNEL32DLL) == 0) {
			iskernel32 = true;
			patch = HookMethod(lpProcName, &saved, kernel32Patch, kernel32Methods, kernel32Save, kernel32Count);
			if (!patch) {
				patch = HookMethod(lpProcName, &saved, Kernel32Patch, Kernel32Methods, Kernel32Save, Kernel32Count);
			}
		} else if (lstrcmpi(nameBuffer, SHELL32DLL) == 0) {
			patch = HookMethod(lpProcName, &saved, shell32Patch, shell32Methods, shell32Save, shell32Count);
		} else if (lstrcmpi(nameBuffer, PROPSYSDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, propsysPatch, propsysMethods, propsysSave, propsysCount);
		} else if (lstrcmpi(nameBuffer, OLE32DLL) == 0) {
			patch = HookMethod(lpProcName, &saved, ole32Patch, ole32Methods, ole32Save, ole32Count);
		} else if (lstrcmpi(nameBuffer, WMVCOREDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, wmvcorePatch, wmvcoreMethods, wmvcoreSave, wmvcoreCount);
		} else if (lstrcmpi(nameBuffer, MFDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, mfPatch, mfMethods, mfSave, mfCount);
		} else if (lstrcmpi(nameBuffer, QUARTZDLL) == 0) {
			patch = HookMethod(lpProcName, &saved, quartzPatch, classObjectMethods, quartzSave, classObjectCount, &method, quartzExternals);
		}
	} catch (...) {
		TraceException(_T("NewGetProcAddress"));
	}
	void* proc = (*(PGetProcAddress)kernel32Save[1])(hModule, lpProcName);
	if (!patch || !saved) {
		patch = proc;
	} else {
		*saved = proc;
		if (!iskernel32 && proc)
			SetPrevFuncAddress(method, proc);
	}
	return (FARPROC)patch;
}

FARPROC WINAPI GetHookedProcAddress(HMODULE module, LPCSTR method) {
	if (!kernel32Save[1])
		kernel32Save[1] = GetProcAddress;
	return GetProcAddressHook(module, method);
}

HMODULE WINAPI GetModuleHandleAHook(LPCSTR lpModuleName) {
	HMODULE patch = (*(PGetModuleHandleA)kernel32Save[2])(lpModuleName);
	HookHandle(&patch);
	return patch;
}

HMODULE WINAPI GetModuleHandleWHook(LPCWSTR lpModuleName) {
	HMODULE patch = (*(PGetModuleHandleW)kernel32Save[3])(lpModuleName);
	HookHandle(&patch);
	return patch;
}

